import InstructionTable, HistoryFile
import utility, os, string, threading, pickle
import scipy, numpy
from Crypto.Util import number
from Crypto.Hash import HMAC, SHA256

"""
To make SecLogin threadsafe:

1. Maintain an array of locks in an associated array, by key being name of user. Also maintain one lock for whole array.
2. Steps of login:
    - check if mutex in array by key username exists. (if(mutex_array['username'])
    - if yes, lock mutex in array by key username. do login, then unlock mutex.
    - if no, then user does not exist, abort.
3. Steps of add:
    - check if mutex in array by key username exists. (if(mutex_array['username'])
    - if yes, then user already exists, cant add, abort.
    - if no, then lock whole array. Here, check again if it doesn't exist -> maybe somebody created while you were trying to lock it.
        if still doesn't exist, add one user mutex to it, then unlock array. then lock user mutex. if it already exists, unlock and abort (it was created while we were waiting).
    - now the corresponding mutex is locked, so execute your add_user command.

SecLogin is threadsafe, SecLogin does not implement threading itself: if there would be multithreading, that would be the interface that would call
SecLogin.add_user, SecLogin.login_user in threads for same SecLogin class instance.
"""
class SecLogin:
    def __init__(self, home, min_pwd_length = 6, max_pwd_length = 20, history_length = None):
        """
        All i/o access right tests are done with except instead of access, to avoid TOCTOU vulnerabilities.
        Note that IOError can mean that rights are denied, not that its not found,
        but we return this msg anyway, a la RHEL.
        """
        #Set home directory for SecLogin.
        #Contains serialized InstructionTables and HistoryFiles, in directorires named by username.
        #also contains config file
        try:
            self.main_lock = threading.Lock()
            self.lock_map = dict()
            
            #abspath() takes care of canonicalization, decoding. now we can check its the right thing.            
            home = os.path.abspath(home)
            
            #currently we allow the SecLogin root directory to be anywhere in the sytem that uid allows to r/w.
            #this could be changed, e.g. be an install time parameter when installing as module.

            self.home = home
            self.len = 160L
            os.chdir(self.home)           
            try:
                cfg_file = open("seclogin.cfg", 'rb')
                utility.debug('Opened config file.')
                #if its successful, then unpickle it to get values
                try:
                    self.min_pwd_length, self.max_pwd_length, self.history_length, self.mod = pickle.load(cfg_file)

                    try:
                        self.min_pwd_length = abs(int(self.min_pwd_length))
                        assert(self.min_pwd_length >= 6)
                        self.max_pwd_length = abs(int(self.max_pwd_length))
                        assert(self.max_pwd_length <= 20)
                        self.history_length = abs(int(self.history_length))
                        assert(self.history_length > 0)
                        self.mod = abs(int(self.mod))
                        assert(number.isPrime(self.mod))
                    except:
                        utility.debug('Invalid values in config file.')
                        return None
                
                    utility.debug('Config file deserialized.')
                except:
                    utility.debug('Unpickling failed.')
                    return None
                                
                
                try:
                    cfg_file.close()
                except IOError:
                    utility.debug('Config file close error.')
                    return None                
                #initialization from config file is done.                
                return                
            except IOError:
                utility.debug('No config file found. SecLogin initialization failed.')
                return None            
        except OSError:
            #utility.debug('Home directory not found.')
            if (history_length == None):
                #its an error, supposed to be a load!
                utility.debug('SecLogin initialization failed.')
                return None
            else:
                try:
                    os.mkdir(self.home)
                    os.chdir(self.home)
                    utility.debug('SecLogin home created.')
                except OSError:
                    utility.debug('Home directory creation failed.')
                    return None
                try:
                    cfg_file = open("seclogin.cfg", 'wb')
                except IOError:
                    utility.debug('No config file found. SecLogin initialization failed.')
                    return None

        try:
            self.min_pwd_length = abs(int(min_pwd_length))
            if (self.min_pwd_length < 6):
                raise ValueError
        except:
            utility.debug('Password length incorrect.')
            try:
                cfg_file.close()
                os.remove(self.home + '/seclogin.cfg')
                os.chdir('../')
                os.rmdir(self.home)
                return None
            except (IOError, OSError):
                return None            
        try:
            self.max_pwd_length = abs(int(max_pwd_length))
            if (self.max_pwd_length > 20):
                raise ValueError
        except:
            utility.debug('Password length incorrect.')
            try:
                cfg_file.close()
                os.remove(self.home + '/seclogin.cfg')
                os.chdir('../')
                os.rmdir(self.home)
                return None
            except (IOError, OSError):
                return None
        
        if (history_length):
            try:
                self.history_length = abs(int(history_length))
            # ValueError, TypeError, AttributeError
            except:
                utility.debug('provided history length not appropriate')
                try:
                    cfg_file.close()
                    os.remove(self.home + '/seclogin.cfg')
                    os.chdir('../')
                    os.rmdir(self.home)
                    return None
                except (OSError, IOError):
                    return None

        #this is fixed so that we dont have to get the len of mod every time.
        #if it is not fixed, it needs to be asserted to be 0 mod 8 here at creation!
        self.len = 160L
        
        
        #check prime functions return
        #q of Zq
        self.mod = self._generate_prime(self.len)
        if (self.mod == None):
            utility.debug('Prime number generation failed.')
            try:
                cfg_file.close()
                os.remove(self.home + '/seclogin.cfg')
                os.chdir('../')
                os.rmdir(self.home)
                return None
            except (OSError, IOError):
                return None
        #deprecated
        self.r = 0

        #pickle the created values into the config file if we are initing, not leading.
        try:                
            pickle.dump([self.min_pwd_length, self.max_pwd_length, self.history_length, self.mod], cfg_file)
        except:
            utility.debug('Pickling failed.')
            cfg_file.close()
            os.remove(self.home + '/seclogin.cfg')

            os.chdir('../')
            os.rmdir(self.home)            
            return None

        try:
            cfg_file.close()
        except IOError:
            utility.debug('Config file close error.')
            cfg_file.close()
            os.remove(self.home + '/seclogin.cfg')
            
            os.chdir('../')
            os.rmdir(self.home)
            return None        

    #Private functions. Implement the mathematical operations of initialization and authentication.

    """
    Check password supplied against password policy. Currently that is:
    - only a-zA-Z
    - min is min_pwd_len, max is max_pwd_len
    """
    def _pwd_policy_check(self, pwd):        
        return (self.min_pwd_length <= len(pwd) and len(pwd) <= self.max_pwd_length and pwd.isalpha())

    """
    Return the case of a character according to SecLogin's case rules.

    Parameters:
    char: the character.

    Output: 0 for lowercase, 1 for uppercase, 0 for numbers and printable special characters
    """
    def _get_case(self, ch):                
        #smaller ascii then 97 means uppercase. that will be True, int(True) is 1, int(False) is 0.
        return int((ord(ch)) < 97)

    """
    Compute feature_vector from password value.

    Parameters:
    pwd: the password.

    Output: the feature vector
    """
    def _get_feature_vector(self, pwd):
        return [self._get_case(i) for i in pwd]
    
    """
    Keyed hash function. Uses Crypto library.

    Parameters:
    pwd: the key for the hash function
    value: the plaintext to generate a digest for

    Output: the hash value
    """

    def _hash(self, pwd, plain):    
        hmac = HMAC.new(pwd, number.long_to_bytes(plain), SHA256)
        return number.bytes_to_long(hmac.digest())

    """
    Generate a prime number of length len. Crypto.Util.isPrime uses Rabin-Miller test.

    Parameters:
    len: length of the prime

    Output: 160 bit random prime.
    """
    def _generate_prime(self, length):
        value = 0        
        try:
            assert(length % 8 == 0)
        except AssertionError:
            utility.debug('Bit length not 8 multiple')
            return None
        while(1):
            rand_bytes = os.urandom(length / 8)
            value = number.bytes_to_long(rand_bytes)
            # Ensure high bit is set: multiply (=bitwise or) it by 2^(len-1)
            value |= 2L ** (length-1)
            count = 0
            while (not number.isPrime(value)):
                value = value + 2
                count = count + 1
                if (count == 1000):
                    break
            if (number.isPrime(value)):
                break
            else:
                pass
        return value

    """
    Generates the random parameter r of len bits. len has to be: len % 8 = 0
    (otherwiser rand_bytes needs to be fitted before transforming we alwaya use it with 160 bits here).
    This function can be used to generate:
    - the q of Zq,
    - parameter r,
    - by generate_polynomial to generate coefficients for the polynomial.

    Parameters:
    len: the length of the parameter in bits
    test: if supplied, make sure that generated number is a different one

    Output: the generated parameter.
    """
    def _generate_param (self, length, test=None):
        value = 0        
        while (1):
            try:
                assert(length % 8 == 0)
            except AssertionError:
                utility.debug('Bit length not 8 multiple')
                return None
            rand_bytes = os.urandom(length / 8)
            value = number.bytes_to_long(rand_bytes)
            # Ensure high bit is set: multiply (=bitwise or) it by 2^(len-1)
            value |= 2L ** (length-1)
            if (value != test):
                break
        return value

    """
    Generates a random polynomial of degree m-1 on Zq

    Parameters:
    hpwd: the value f(0) to be used for the generated polynomial
    degree: the degree of the polynomial

    Output: the generated polynomial, as poly1d value
    """
    def _generate_polynomial(self, hpwd, degree):        
        
        #f[x] = a*x^2 + b*x + c has degree two, f(0) = c, so we need degree number of coefficients
        
        coeffs = [self._generate_param(self.len) % self.mod for i in range(degree)]
        coeffs.append(hpwd)

        return numpy.poly1d(coeffs)

    """
    Generates a random point on the polynomial pol

    Parameters:
    i: the iterator used for generation of the x value of the point
    pol: the polynomial.

    Output: the generated y0 and y1 coordinates of the point x on the polynomial. (we know x is 2i+1, so no point in passing it around.)
    This function is used by _compute_y_pairs().
    """

    def _generate_point(self, i, pol):
        #return (pol(2*i), pol(2*i + 1))        
        return (pol(2*i) % self.mod, pol(2*i + 1) % self.mod)

    """
    Computes the array of y_a0, y_a1 pairs.

    Parameters:
    feature_vector: the feature vector for the users password.
    f[i] = -1, if lower_case only for length h in history file, 1 if upper_case only, 0 if either
    pol: the random polynomial used

    Output: the computed y_a0, y_a1 pairs.
    """

    def _compute_y_pairs(self, feature_vector, pol):        
        feature_vector_length = len(feature_vector)
        try:
            assert(feature_vector_length > 0)
        except AssertionError:
            utility.debug('Password length is not suitable.')

        #for each feature, get x, compute its y, then return the pair (putting garbage where needed)
        #Note, that the i value is +1'ed, so that i goes from 1 to m, not 0 to m-1 (paper)
        xvals = [i for i in range(1, feature_vector_length + 1)]
        true_yvals = map(lambda x: self._generate_point(x, pol), xvals)

        y_pairs = []

        for i in range(feature_vector_length):
            #feature distinguishing
            if (abs(feature_vector[i]) == 1):
                #uppercase                
                if(feature_vector[i] == 1):
                    #myt = tuple([ self._generate_param(self.len, true_yvals[i][0]), true_yvals[i][1] ])
                    y_pairs.append(tuple([self._generate_param(self.len, true_yvals[i][0]), true_yvals[i][1]]))
                #lowercase
                else:                    
                    y_pairs.append(tuple([ true_yvals[i][0], (self._generate_param(self.len, true_yvals[i][1]))])) 
            #feature undistinguishing
            else:
                y_pairs.append(tuple([true_yvals[i][0], true_yvals[i][1]]))            
        return y_pairs

    """
    Generates the instruction table.
    It calculates an alpha, beta value for each row.

    Parameters:
    pwd: the password of the user
    y_pairs: array of pairs of y_a0, y_a1 values.

    This is basically the contructor for an Instruction Table, more preciselly: the output of it can be used as such.

    Output: the generated instruction table, as a two-dimensional array.
    """

    def _compute_instruction_table(self, pwd, y_pairs):
        instruction_table_rows = []
        
        for i in range(len(y_pairs)):
            # i goes from 0 to m-1, but it is 1 to m according to paper, so +1
            
            alpha = (y_pairs[i][0] + self._hash(pwd.lower(), 2*(i+1))) % self.mod
            beta = (y_pairs[i][1] + self._hash(pwd.lower(), 2*(i+1) + 1)) % self.mod            
            instruction_table_rows.append(tuple([alpha, beta]))
            
        return instruction_table_rows

    """
    Computes the points on the polynomial pol from pwd, alpha, beta.

    Parameters:
    pwd: the users password, used in G() and to determine which column to use
    from the instruction table.
    table: the instruction table used here to recover points: the actual InstructionTable object
    Note: in x calculation, we use i+1 instead of i because i value runs from 1 to m, not 0 to m-1.
    Every other usage of i is as index into the list (not the actual value), and internally everything
    is indexed from 0 to m-1.
    
    Output: the list of computed (x,y) coordinates of the point on the polynomial
    """

    def _compute_points(self, pwd, table):
        points = []
        #use as lowercase for encryption/decryption
        pwd_encrypt = pwd.lower()
        for i in range(len(pwd)):
            if (self._get_case(pwd[i]) == 0):
                x = 2*(i+1)
                alpha = table.get_alpha(i, pwd_encrypt)            
                
                if (alpha == None):
                    utility.debug('Alpha recovery failed!')
                    return None
                y = (alpha - self._hash(pwd_encrypt, 2*(i+1))) % self.mod
            else:
                x = (2*(i+1) + 1)
                beta = table.get_beta(i, pwd_encrypt)
                if (beta == None):
                    utility.debug('Beta recovery failed!')
                    return None                
                y = (beta - self._hash(pwd_encrypt, 2*(i+1) + 1)) % self.mod
            points.append(tuple([x,y]))

        return points

    """
    polynomials = []                    
    for i in pairs.keys():        
        _num = reduce(lambda x,y: x*y, map(lambda x: scipy.poly1d([1, -x]), [j for j in pairs.keys() if i != j]))
        _denum = reduce(lambda x,y: x*y, map(lambda x: x[0] - x[1], [[i,j] for j in pairs.keys() if i != j]))        
        polynomials.append((_num/_denum) * pairs.get(i))
    return reduce(lambda x,y: x+y, polynomials)
    """
                
    """
    Computes the lagrange coefficient lambda.

    Parameters:
    coeffs: list of the computed x_i values

    Output: the list of lagrange_coefficients
    """
    def _compute_lagrange_coeffs(self, coeffs):
        lgr_coeffs = []
        for i in range(len(coeffs)):                        
            num = reduce(lambda x,y: x*y % self.mod, map(lambda j: coeffs[j], [j for j in range(len(coeffs)) if j != i]))
            denum = reduce(lambda x,y: x*y % self.mod, map(lambda j: coeffs[j] - coeffs[i], [j for j in range(len(coeffs)) if j != i]))
            lgr_coeffs.append(number.inverse(denum, self.mod)*num % self.mod)
        return lgr_coeffs
            
    """
    Computes the hpwd.

    Parameters:
    lgr_coeffs: list of the lagrange coefficients
    points: list of the y_i values for the points on the polynomial

    Output: the calculated hwpd hardenned password.

    hpwd = SUM(lgr_coeff_i * y_i) mod q
    """
    def _compute_hpwd(self, lgr_coeffs, points):

        sum_terms = map(lambda x,y: (x*y) % self.mod, lgr_coeffs, points)
        return reduce(lambda x,y: (x+y) % self.mod, sum_terms)

    #Public functions.

    def get_min_pwd_length(self):
        return self.min_pwd_length

    def new_user(self, username, pwd):

        #test the username and pwd are strings.
        try:
            assert(type(username) == str)
            assert(type(pwd) == str)
        except:
            utility.debug('Invalid parameters.')
            return False

        utility.debug('SecLogin: adding new user.')

        try:
            os.chdir(self.home)
        except OSError:
            utility.debug('SecLogin: Home directory not found.')
            return False        

        dirname = os.path.normpath(username)
        #Make sure its a library name: does not contain path traversal characters.
        if ((dirname.find('\\') != -1) or (dirname.find('/') != -1)):
            utility.debug('Invalid path.')
            return False
        
        path = os.path.abspath(username)
        #must be: self.home/username
        if (os.path.dirname(path) != self.home or os.path.basename(path) != username):
            utility.debug('Invalid path.')
            return False        

        if (self._pwd_policy_check(pwd) == False):
            utility.debug('Password does not match policy.')
            return False

        #acquire lock        
        if (self.lock_map.get(username) != None):
            utility.debug('This username has already been initialized.')
            return False
        else:
            self.main_lock.acquire()
            if (self.lock_map.get(username) != None):
                utility.debug('This username has already been initialized.')
                self.main_lock.release()
                return False
            else:
                new_lock = {username : threading.Lock()}
                self.lock_map.update(new_lock)
                self.lock_map[username].acquire()                

        #init the insTable from path
        insTable = InstructionTable.InstructionTable(pwd, path, True)

        if (insTable == None):
            utility.debug('Instruction Table creation failed.')
            self.main_lock.release()
            self.lock_map[username].release() 
            return False

        try:    
            assert(insTable.is_serialized() == False)
        except AssertionError:
            utility.debug('User already exists, adding failed.')
            self.main_lock.release()
            self.lock_map[username].release()
            return False

        #new random polynomial -> new y pairs -> new alpha, beta -> updated insTable
        hpwd = self._generate_param(self.len) % self.mod

        utility.debug('SecLogin: originally calculated pwd is %d' % hpwd)
        
        pol = self._generate_polynomial(hpwd, len(pwd) - 1)
        #Original feature vector is all zeros: until we don't have h inputs, all is accepted.
        if (self.history_length > 1):
            feature_vector = [0 for i in range(len(pwd))]
        #the exception is history_length = 1, which is essentially the no hardening case (talk about much ado about nothing :)
        else:
            feature_vector = self._get_feature_vector(pwd)
        y_pairs = self._compute_y_pairs(feature_vector, pol)
        
        rows = self._compute_instruction_table(pwd, y_pairs)

        if (insTable.update_rows(rows, pwd.lower()) == False):
            utiltiy.debug('Instruction Table update failed.')
            self.main_lock.release()
            self.lock_map[username].release()
            return False            
        if (insTable.serialize() == False):
            utiltiy.debug('Instruction Table serialization failed.')
            self.main_lock.release()
            self.lock_map[username].release()
            return False                        
        
        hFile = HistoryFile.HistoryFile(len(pwd), self.history_length, path, hpwd)

        if (hFile == None):
            utiltiy.debug('History File initialization failed.')
            self.main_lock.release()
            self.lock_map[username].release()
            return False

        #update and serialize with hpwd.
        if ((hFile.update(self._get_feature_vector(pwd))) == False):
            #can not update yet, because it is not deserialized, this should never happen right after init, btw
            if (hFile.deserialize() == False):
                utility.debug('Deserialization failed.')
                self.main_lock.release()
                self.lock_map[username].release()
                return False
            if (hFile.update(self._get_feature_vector(pwd)) == False):
                #this should NEVER happen if deserialize() is successful.
                utility.debug('History File update failed.')
                self.main_lock.release()
                self.lock_map[username].release()
                return False
        #serialize HistoryFile, now with new hpwd.
        
        if (hFile.serialize() == False):
            utility.debug('History File serialize failed while adding user.')                
            self.main_lock.release()
            self.lock_map[username].release()
            return False        
                
        utility.debug('Successfully added new user.')
        self.main_lock.release()
        self.lock_map[username].release()
        return

    def login_user(self, username, pwd):

        #test the username and pwd are strings.
        try:
            assert(type(username) == str)
            assert(type(pwd) == str)
            
        except:
            utility.debug('Invalid parameters.')
            return False

        if (self._pwd_policy_check(pwd) == False):
            utility.debug('Password does not match policy.')
            return False

        utility.debug('SecLogin: log user in.')
        #get insTable path from username
        
        #Make sure we are in the right directory.
        try:
            os.chdir(self.home)
        except OSError:
            utility.debug('Home directory not found.')
            return False
        
        path = os.path.abspath(username)
        #its properly canonicalized, decoded now. make sure its the right path.
        #must be: self.home/username, username must not contain /

        dirname = os.path.normpath(username)
        #Make sure its a library name: does not contain path traversal characters.
        if ((dirname.find('\\') != -1) or (dirname.find('/') != -1)):
            utility.debug('Invalid path.')
            return False
        
        #must be: self.home/username
        if (os.path.dirname(path) != self.home or os.path.basename(path) != username):
            utility.debug('Invalid path.')
            return False

        #acquire lock
        if (self.lock_map.get(username) == None):
            utility.debug('This user does not exist.')
            return False
        else:
            try:
                self.lock_map[username].acquire()
            except:
                utility.debug('Acquiring lock failed.')
                return False

        #init the insTable from path
        insTable = InstructionTable.InstructionTable(pwd, path, False)

        if (insTable == None):
            utility.debug('Instruction Table creation failed.')
            self.lock_map[username].release()
            return False
        
        if ((insTable.deserialize()) == False):
           utility.debug('Authentication Failed.')
           self.lock_map[username].release()
           return False

        #returns (x,y) pairs
        points = self._compute_points(pwd, insTable)

        if (points == None):
            utility.debug('x,y computation from Instruction Table failed.')
            self.lock_map[username].release()
            return False

        #get xs from points as coeffs
        x_values = [x[0] for x in points]
        
        #returns lgr_coeffs from the xs from (x,y) pairs
        lgr_coeffs = self._compute_lagrange_coeffs(x_values)

        #get ys from points as points
        y_values = [x[1] for x in points]
        
        #returns hpwd from ys and lgr_coeffs
        hpwd = self._compute_hpwd(lgr_coeffs, y_values)

        utility.debug('SecLogin: calculated pwd is %d' % hpwd)

        #init (decrypt and dserialize) HF. If this is successful, authentication is OK.
        hFile = HistoryFile.HistoryFile(len(pwd), self.history_length, path, hpwd)

        if (hFile == None):
            utilty.debug('History File initialization error, login failed.')
            self.lock_map[username].release()
            return False
        else:
            utility.debug('History File initialized and deserialized.')
        
        #update HistoryFile.
        if (hFile.update(self._get_feature_vector(pwd)) == False):
            #this should NEVER happen if deserialize() is successful.
            utility.debug('Can not update, it is not deserialized.')
            self.lock_map[username].release()
            return False
        #serialize HistoryFile, now with new hpwd.
        feature_vector = hFile.get_feature_vector()
        if (feature_vector == False):
            utility.debug('Getting feature vector failed.')
            self.lock_map[username].release()
            return False     
        if (hFile.serialize() == False):
            utility.debug('Serialize failed.')
            self.lock_map[username].release()
            return False        
    
        #new random polynomial -> new y pairs -> new alpha, beta -> updated insTable
        pol = self._generate_polynomial(hpwd, len(pwd) - 1)   
        y_pairs = self._compute_y_pairs(feature_vector, pol)
        rows = self._compute_instruction_table(pwd, y_pairs)

        if (insTable.update_rows(rows, pwd.lower()) == False):
            utility.debug('Instruction table update failed.')
            self.lock_map[username].release()
            return False            
     
        if (insTable.serialize() == False):
            utility.debug('Instruction table serialization failed.')
            self.lock_map[username].release()
            return False

        utility.debug('SecLogin: login successful. Welcome!')
        self.lock_map[username].release()
        return

#Insert unit test here for verifying module implementation
if __name__ == "__main__":
    print 'Unit Test'
    password = 'abcdef'
    sL = SecLogin('/tmp/test_sl', 6, 9, 4)
    assert(sL != None)
    sL.new_user('myname', 'abcdef')
    #History length is 4, all is accepted
    try:
        assert(sL.login_user('myname', 'abcdeF') != False)
        assert(sL.login_user('myname', 'abcDEf') != False)
        assert(sL.login_user('myname', 'ABCDEf') != False)
        assert(sL.login_user('myname', 'abcdef') != False)
        #now fill in with same
        assert(sL.login_user('myname', 'abcDEF') != False)
        assert(sL.login_user('myname', 'abcDEF') != False)
        assert(sL.login_user('myname', 'abcDEF') != False)
        #now only [-1,-1,-1,0,0,0], but last 3 are the same ->
        #   one more makes history_file lock in at one case value.
        assert(sL.login_user('myname', 'abcDEF') != False)
        #now only [-1,-1,-1,1,1,1] -> abcDEF accepted but nothing else
        assert(sL.login_user('myname', 'abcdef') == False)
        assert(sL.login_user('myname', 'ABCDEF') == False)
        assert(sL.login_user('myname', 'abcDEF') != False)
        print 'Test finished.'
    except AssertionError:
        print 'Unit Test Failed.'
